#大数据2201刘源媛
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import gridspec
import tushare as ts
from matplotlib.dates import date2num
import datetime
from tqdm import tqdm
import matplotlib as mpl

# 设置中文显示问题
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
mpl.rcParams['font.size'] = 12  # 设置字体大小

# 设置Tushare token
ts.set_token('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
pro = ts.pro_api()


class StockTradingStrategy:
    def __init__(self):
        # 策略参数配置
        self.config = {
            # 大阳线定义：涨幅超过5%且成交量较前一日放大50%
            'BIG_UP_PCT': 0.05,
            'VOLUME_INCREASE': 0.5,

            # 选股参数
            'HOLD_DAYS': 5,  # 持有天数
            'STOP_LOSS': -0.05,  # 止损比例
            'STOP_GAIN': 0.10,  # 止盈比例

            # 回测参数
            'START_DATE': '20200101',
            'END_DATE': '20201231',
            'INIT_CAPITAL': 100000,  # 初始资金
            'COMMISSION': 0.0003,  # 交易佣金
            'MAX_STOCKS': 10  # 最大选股数量
        }

        # 获取股票基础数据
        self.stock_basic = self.get_stock_basic()

    def get_stock_basic(self):
        """获取股票基础信息"""
        df = pro.stock_basic(exchange='', list_status='L',
                             fields='ts_code,symbol,name,industry,list_date')
        return df

    def get_daily_data(self, ts_code, start_date, end_date):
        """获取日线数据"""
        df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
        df = df.sort_values('trade_date')
        df['trade_date'] = pd.to_datetime(df['trade_date'])
        df.set_index('trade_date', inplace=True)
        return df

    def plot_candlestick(self, ax, df, width=0.6):
        """手动绘制K线图"""
        width = width / 24 / 60 / 60  # 转换为天的分数

        for idx, row in df.iterrows():
            date = date2num(idx)
            open_price = row['open']
            close_price = row['close']
            high_price = row['high']
            low_price = row['low']

            if close_price >= open_price:
                color = 'red'
                lower = open_price
                upper = close_price
            else:
                color = 'green'
                lower = close_price
                upper = open_price

            ax.plot([date, date], [low_price, high_price], color=color, linewidth=1)
            rect = plt.Rectangle((date - width / 2, lower), width, upper - lower,
                                 facecolor=color, edgecolor=color)
            ax.add_patch(rect)

        ax.xaxis_date()
        ax.autoscale_view()

    def generate_signals(self, stock_data):
        """生成交易信号"""
        df = stock_data.copy()

        # 计算涨跌幅和成交量变化
        df['pct_chg'] = df['close'].pct_change()
        df['vol_chg'] = df['vol'].pct_change()

        # 生成买入信号（大阳线）
        df['buy_signal'] = (df['pct_chg'] > self.config['BIG_UP_PCT']) & \
                           (df['vol_chg'] > self.config['VOLUME_INCREASE'])

        # 初始化卖出信号
        df['sell_signal'] = False

        # 对于每个买入信号，计算后续的卖出条件
        buy_indices = df[df['buy_signal']].index
        for buy_date in buy_indices:
            hold_period = df.loc[buy_date:].iloc[:self.config['HOLD_DAYS']]
            hold_period['cum_pct'] = hold_period['close'] / hold_period.iloc[0]['close'] - 1

            stop_condition = (hold_period['cum_pct'] < self.config['STOP_LOSS']) | \
                             (hold_period['cum_pct'] > self.config['STOP_GAIN'])

            if stop_condition.any():
                first_sell_date = hold_period[stop_condition].index[0]
                df.loc[first_sell_date, 'sell_signal'] = True
            else:
                last_date = hold_period.index[-1]
                df.loc[last_date, 'sell_signal'] = True

        return df

    def select_stocks(self):
        """选股逻辑"""
        # 只保留A股主板股票
        stocks = self.stock_basic[~self.stock_basic['ts_code'].str.contains('BJ')]
        results = []

        print("\n正在筛选符合条件的股票...")
        for ts_code in tqdm(stocks['ts_code'].sample(min(100, len(stocks))), desc="筛选股票"):  # 随机选取最多100只股票进行筛选
            try:
                data = self.get_daily_data(ts_code, self.config['START_DATE'], self.config['END_DATE'])
                if len(data) < 30:  # 忽略交易天数太少的股票
                    continue

                # 计算大阳线
                data['pct_chg'] = data['close'].pct_change()
                data['vol_chg'] = data['vol'].pct_change()
                big_up_days = data[(data['pct_chg'] > self.config['BIG_UP_PCT']) &
                                   (data['vol_chg'] > self.config['VOLUME_INCREASE'])]

                if len(big_up_days) > 0:
                    avg_pct = big_up_days['pct_chg'].mean()
                    results.append({
                        'ts_code': ts_code,
                        'name': stocks[stocks['ts_code'] == ts_code]['name'].values[0],
                        'big_up_count': len(big_up_days),
                        'avg_pct': avg_pct,
                        'last_close': data.iloc[-1]['close']
                    })
            except Exception as e:
                continue

        # 转换为DataFrame并排序
        results_df = pd.DataFrame(results)
        if len(results_df) > 0:
            results_df['score'] = results_df['big_up_count'] * results_df['avg_pct']
            results_df = results_df.sort_values('score', ascending=False)
            results_df = results_df.head(self.config['MAX_STOCKS'])  # 只保留前N只股票

        return results_df

    def backtest(self, ts_code):
        """回测单个股票的策略表现"""
        data = self.get_daily_data(ts_code, self.config['START_DATE'], self.config['END_DATE'])
        data = self.generate_signals(data)

        # 初始化回测变量
        capital = self.config['INIT_CAPITAL']
        position = 0
        trades = []
        buy_price = 0
        hold_days = 0
        daily_values = []

        # 遍历每一天
        for i, row in data.iterrows():
            # 如果有持仓，更新持有天数
            if position > 0:
                hold_days += 1

            # 卖出信号
            if row['sell_signal'] and position > 0:
                sell_price = row['close']
                pct_chg = (sell_price - buy_price) / buy_price
                profit = position * pct_chg
                capital += position * (1 + pct_chg) * (1 - self.config['COMMISSION'])

                trades.append({
                    'date': i,
                    'type': 'sell',
                    'price': sell_price,
                    'shares': position,
                    'profit': profit,
                    'hold_days': hold_days
                })

                position = 0
                buy_price = 0
                hold_days = 0

            # 买入信号
            if row['buy_signal'] and position == 0 and capital > 0:
                buy_price = row['close']
                position = capital * (1 - self.config['COMMISSION']) / buy_price
                capital = 0

                trades.append({
                    'date': i,
                    'type': 'buy',
                    'price': buy_price,
                    'shares': position
                })

            # 记录每日价值
            if position > 0:
                daily_value = capital + position * row['close'] * (1 - self.config['COMMISSION'])
            else:
                daily_value = capital
            daily_values.append(daily_value)

        # 计算回测结果
        final_value = capital + position * data.iloc[-1]['close'] * (1 - self.config['COMMISSION'])
        total_return = (final_value - self.config['INIT_CAPITAL']) / self.config['INIT_CAPITAL']

        # 计算年化收益率
        days = (data.index[-1] - data.index[0]).days
        annualized_return = (1 + total_return) ** (365 / days) - 1 if days > 0 else 0

        # 计算最大回撤
        data['portfolio_value'] = daily_values
        data['cummax'] = data['portfolio_value'].cummax()
        data['drawdown'] = (data['portfolio_value'] - data['cummax']) / data['cummax']
        max_drawdown = data['drawdown'].min()

        # 计算胜率
        if len(trades) > 0:
            winning_trades = [t for t in trades if t.get('profit', 0) > 0]
            win_rate = len(winning_trades) / len([t for t in trades if 'profit' in t])
        else:
            win_rate = 0

        return {
            'ts_code': ts_code,
            'name': self.stock_basic[self.stock_basic['ts_code'] == ts_code]['name'].values[0],
            'final_value': final_value,
            'total_return': total_return,
            'annualized_return': annualized_return,
            'max_drawdown': max_drawdown,
            'win_rate': win_rate,
            'trades': trades,
            'data': data
        }

    def visualize_strategy(self, result):
        """可视化策略表现"""
        data = result['data'].copy()
        trades = result['trades']

        # 创建图表
        fig = plt.figure(figsize=(14, 12))
        gs = gridspec.GridSpec(4, 1, height_ratios=[3, 1, 1, 1])

        # K线图
        ax1 = plt.subplot(gs[0])
        self.plot_candlestick(ax1, data)

        # 标记买卖点
        buy_dates = [t['date'] for t in trades if t['type'] == 'buy']
        buy_prices = [t['price'] for t in trades if t['type'] == 'buy']
        sell_dates = [t['date'] for t in trades if t['type'] == 'sell']
        sell_prices = [t['price'] for t in trades if t['type'] == 'sell']

        ax1.plot(buy_dates, buy_prices, '^', markersize=10, color='m', label='买入')
        ax1.plot(sell_dates, sell_prices, 'v', markersize=10, color='k', label='卖出')

        # 设置标题和标签
        title = f"{result['name']} ({result['ts_code']}) - 大阳线策略\n" \
                f"总收益率: {result['total_return'] * 100:.2f}% | " \
                f"年化收益率: {result['annualized_return'] * 100:.2f}% | " \
                f"最大回撤: {result['max_drawdown'] * 100:.2f}% | " \
                f"胜率: {result['win_rate'] * 100:.2f}%"
        ax1.set_title(title, fontproperties='SimHei')
        ax1.set_ylabel('价格')
        ax1.grid(True)
        ax1.legend(prop={'family': 'SimHei'})

        # 成交量
        ax2 = plt.subplot(gs[1], sharex=ax1)
        colors = ['r' if close >= open else 'g' for close, open in zip(data['close'], data['open'])]
        ax2.bar(data.index, data['vol'], color=colors)
        ax2.set_ylabel('成交量')
        ax2.grid(True)

        # 资金曲线
        ax3 = plt.subplot(gs[2], sharex=ax1)
        ax3.plot(data.index, data['portfolio_value'], label='组合价值')
        ax3.set_ylabel('价值')
        ax3.grid(True)

        # 回撤曲线
        ax4 = plt.subplot(gs[3], sharex=ax1)
        ax4.fill_between(data.index, data['drawdown'], 0, color='red', alpha=0.3)
        ax4.set_ylabel('回撤率')
        ax4.grid(True)

        plt.gcf().autofmt_xdate()
        plt.tight_layout()
        plt.show()

    def compare_strategies(self, results):
        """比较多只股票的策略表现"""
        plt.figure(figsize=(14, 8))

        # 绘制每条收益率曲线
        for result in results:
            data = result['data']
            plt.plot(data.index, data['portfolio_value'] / self.config['INIT_CAPITAL'],
                     label=f"{result['name']} ({result['total_return'] * 100:.1f}%)")

        # 添加基准线
        plt.axhline(y=1, color='gray', linestyle='--', label='基准')

        plt.title('多股票策略表现比较', fontproperties='SimHei')
        plt.ylabel('累计收益率')
        plt.xlabel('日期')
        plt.grid(True)
        plt.legend(prop={'family': 'SimHei'})
        plt.tight_layout()
        plt.show()

    def run(self):
        """运行策略"""
        # 1. 选股
        print("正在选股...")
        selected_stocks = self.select_stocks()

        if len(selected_stocks) == 0:
            print("没有符合条件的股票")
            return

        print("\n选股结果:")
        print(selected_stocks[['ts_code', 'name', 'big_up_count', 'avg_pct', 'score']])

        # 2. 回测选中的股票
        results = []
        print("\n开始回测...")
        for _, stock in tqdm(selected_stocks.iterrows(), total=len(selected_stocks), desc="回测进度"):
            ts_code = stock['ts_code']
            result = self.backtest(ts_code)
            results.append(result)

            print(f"\n回测结果 ({result['name']}):")
            print(f"初始资金: {self.config['INIT_CAPITAL']:,.2f}")
            print(f"最终价值: {result['final_value']:,.2f}")
            print(f"总收益率: {result['total_return'] * 100:.2f}%")
            print(f"年化收益率: {result['annualized_return'] * 100:.2f}%")
            print(f"最大回撤: {result['max_drawdown'] * 100:.2f}%")
            print(f"交易次数: {len([t for t in result['trades'] if t['type'] == 'buy'])}")
            print(f"胜率: {result['win_rate'] * 100:.2f}%")

            # 3. 可视化
            self.visualize_strategy(result)

        # 4. 比较多股票表现
        self.compare_strategies(results)


if __name__ == "__main__":
    strategy = StockTradingStrategy()
    strategy.run()